import re

import pandas as pd
from joblib import Parallel, delayed
from tqdm import tqdm

from iflytech_assistant.client import instruct_with_profile

df = pd.read_excel(".vscode/程磊标签标注.xlsx")


SYSTEM_PROMPT = """
Role: 聊天助理

## Goals:

- 理解用户的输入内容，对内容进行优化表达，帮助用户更好地实现沟通目标

## Background:

- 用户在即时通信软件中和别人聊天

## Skills:

- 熟悉中文网络流行语
- 擅长理解用户输入背后真正的意图
- 对于用户的问题，能区分用户是向聊天对象提出问题，还是向你提出问题
- 擅长使用各种沟通和语言技巧，把用户简短的输入，扩充优化成更精彩的内容

## Constrains:

- 你不要和用户对话，而是帮助用户优化表达
- 如果用户输入的内容太短，不要向用户寻求更多信息，而是进行适当的扩写
- 生成的内容要适合发送给对方，不要说"嘿"
- **你必须做到：有足够的诚意，生成的内容长一点，起码超过20个字，要和原始输入有较大差异**

## Workflow:

- 不着急，充分理解用户的原始输入，分析用户沟通的意图、沟通的目的，特别是想要表达的情绪
- 慢慢来，充分发挥你的技能，在不改变用户的沟通意图、沟通目的、情绪的前提下，生成更加丰富的表达
"""

USER_PROMPT = """
## 情景描述
我有一句话想对对方说，但是不知道怎么表达。
{tags}

请帮我按照上面的要求，高情商表达以下内容
```
{user_input}
```

## 输出格式

- 每条建议使用一个用户偏好标签风格的句式、风格、语气或人设。
- 输出五条建议后，分别说明每条建议使用了哪种用户偏好标签风格。

输出格式示例（xxx直接是改写结果，不要有解释，标签和额外的内容）

1. xxx
2. xxx
3. xxx
4. xxx
5. xxx

每条建议的解释：
a. xxx
b. xxx
c. xxx
d. xxx
e. xxx
"""


def run_single(line: pd.Series):
    user_input = line["userinput"]

    tags = ""
    # columns = ["user_input", "人设标签", "语气", "风格", "详细标签", "内容标签"]
    if not pd.isna(line["人设标签"]):
        tags += f"改写结果需要体现我的基本人设：{line['人设标签']}\n"
    if not pd.isna(line["语气"]):
        tags += f"可以从这些语气中选择适当的语气进行表达：{line['语气']}\n"
    if not pd.isna(line["风格"]):
        tags += f"改写结果需要体现这些风格：{line['风格']}\n"
    if not pd.isna(line["详细标签"]):
        tags += f"改写结果需要体现这些句式特点：{line['详细标签']}\n"
    if not pd.isna(line["内容标签"]):
        tags += f"你可以结合这些相关内容进行改写：{line['内容标签']}\n"

    user_prompt = USER_PROMPT.format(tags=tags, user_input=user_input)

    response = instruct_with_profile(user_prompt, system_prompt=SYSTEM_PROMPT)
    return response, user_prompt


if __name__ == "__main__":
    result_dict = {
        "system_prompt": [],
        "user_prompt": [],
        "response": [],
    }
    tasks = []
    for i in tqdm(range(len(df))):
        line = df.loc[i]
        tasks.append(delayed(run_single)(line))

    results = Parallel(n_jobs=10, verbose=10)(tasks)

    for response, user_prompt in results:
        result_dict["system_prompt"].append(SYSTEM_PROMPT)
        result_dict["user_prompt"].append(user_prompt)
        result_dict["response"].append(response.llm_response)

    # add these results to the original dataframe
    result_df = pd.DataFrame(result_dict)
    df = pd.concat([df, result_df], axis=1)
    df.to_excel(".vscode/程磊标签标注_改写结果.xlsx", index=False)
